package com.zw.cloud.tools.threadlocal;

import com.alibaba.ttl.TransmittableThreadLocal;
import com.alibaba.ttl.TtlRunnable;
import com.alibaba.ttl.threadpool.TtlExecutors;
import com.zw.cloud.common.thread.pool.ThreadExecutorPool;
import org.apache.commons.collections4.MapUtils;
import org.apache.commons.lang3.StringUtils;

import java.util.HashMap;
import java.util.Map;
import java.util.Objects;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;

/**
 *  TtlRunnable
 *  TtlCallable
 * getTtlExecutor：修饰接口Executor
 * getTtlExecutorService：修饰接口ExecutorService
 * getTtlScheduledExecutorService：修饰接口ScheduledExecutorService
 */
public class TransmittableThreadLocalUtils {
   /* protected static final TransmittableThreadLocal<String> THREAD_LOCAL = new TransmittableThreadLocal<>();
    public static String getValue() {
       return THREAD_LOCAL.get();
    }

    public static void setValue(String value) {
        THREAD_LOCAL.set(value);
    }

    public static void clear() {
        THREAD_LOCAL.remove();
    }*/

    protected static final TransmittableThreadLocal<Map<String,String>> THREAD_LOCAL = new TransmittableThreadLocal<>(){
        protected Map<String, String> initialValue() {
            return (Map<String, String>) new ConcurrentHashMap();
        }
    };
    public static String getValue(String key) {
        if (Objects.isNull(THREAD_LOCAL.get())) {
            return null;
        }
        return (String)((Map)THREAD_LOCAL.get()).get(key);
    }

    public static void setValue(String key, String value) {
        if (Objects.nonNull(THREAD_LOCAL.get()) && Objects.nonNull(key) && Objects.nonNull(value)) {
            ((Map)THREAD_LOCAL.get()).put(key, value);
        }
    }

    public static void clear() {
        THREAD_LOCAL.remove();
    }

    public static void main(String[] args) {
        test1();
        //test2();
    }

    private static void test1() {
        int requestNum = 3;
        int currentRequest = 0;
        ExecutorService ttlExecutorService = TtlExecutors.getTtlExecutorService(ThreadExecutorPool.msgThreadPoolExecutor);
        while (currentRequest++ < requestNum) {
            TransmittableThreadLocalUtils.setValue("key",String.valueOf(currentRequest));
            System.out.println(String.format("处理开始,请求号=%s, value=%s", currentRequest,TransmittableThreadLocalUtils.getValue("key")));

            int finalCurrentRequest = currentRequest;
            // 使用 TtlRunnable 包装; TtlRunnable;TtlCallable
            new Thread(TtlRunnable.get(() -> {
                System.out.println(String.format("异步场景下，请求号=%s, value=%s", finalCurrentRequest, TransmittableThreadLocalUtils.getValue("key")));
            })).start();

            // 使用 TtlExecutors 包装线程池
            /**
             * getTtlExecutor：修饰接口Executor
             * getTtlExecutorService：修饰接口ExecutorService
             * getTtlScheduledExecutorService：修饰接口ScheduledExecutorService
             */
            ttlExecutorService.execute(() -> {
                System.out.println(String.format("异步场景下-pool，请求号=%s, value=%s", finalCurrentRequest, TransmittableThreadLocalUtils.getValue("key")));
            });

            System.out.println(String.format("处理结束,请求号=%s, value=%s", currentRequest, TransmittableThreadLocalUtils.getValue("key")));
            TransmittableThreadLocalUtils.clear();
        }

        System.out.println("========================");
    }
    private static void test2() {
        // 异步处理线程池
        ExecutorService executorService = TtlExecutors.getTtlExecutorService(Executors.newSingleThreadExecutor());
        int requestNum = 3;
        int currentRequest = 0;
        while (currentRequest++ < requestNum) {
            // 请求前初始化traceContext
            TransmittableThreadLocalUtils.setValue("key",String.valueOf(currentRequest));
            System.out.println(String.format("处理开始,请求号=%s, value=%s", currentRequest,TransmittableThreadLocalUtils.getValue("key")));

            // 异步处理，传递traceContext
            int finalCurrentRequest = currentRequest;
            executorService.submit(() -> {
                System.out.println(String.format("异步场景下，请求号=%s, value=%s", finalCurrentRequest, TransmittableThreadLocalUtils.getValue("key")));
            });

            // 请求处理结束清理traceContext
            System.out.println(String.format("处理结束,请求号=%s, value=%s", currentRequest, TransmittableThreadLocalUtils.getValue("key")));
            TransmittableThreadLocalUtils.clear();
        }
    }
}
